# 此脚本使用 Ultralytics 库中的 auto_annotate 函数进行自动标注任务。
# 自动标注是指利用目标检测模型det_model和分割模型sam_model对指定数据集data中的图像进行标注生成。
# 生成结果为segment格式。如果要转化为box格式，还需要做处理。
from ultralytics.data.annotator import auto_annotate
auto_annotate(
    data="/root/autodl-tmp/yolov8/datasets_balanced/valid/images", 
    det_model="/root/autodl-tmp/yolov8/runs/detect/train12/weights/best.pt", # default: det_model="yolov8x.pt"
    sam_model="sam2_t.pt",     # default:  # sam_model="sam2_b.pt"
    device="cuda",
    conf=0.2,
    iou=0.2,
    classes=[0,1,2])